﻿namespace SpikingAnalyticsLib

open System
open System.IO
open System.Collections.Generic
open SpikingNeuronLib
open SpikingAnalyticsLib.GeneralUtilities
open SpikingAnalyticsLib.PatternExtensions

module internal PatternUtilities =

    // Generate a firing pattern from a sequence of (time, neuron) pairs
    let GetPattern pattern inputPeriod =
        let MakePattern firingLists =
            let input = new List<FiringEvent>(Seq.length firingLists)
            firingLists
            |> Seq.iter (fun firingList -> input.AddRange(firingList))
            Some(new Pattern(input))
        pattern
        |> Seq.map (fun (time, neuron) ->
            [for i in 0..inputPeriod..999 do yield FiringEvent(i + time, neuron, EventLabel.Unspecified)]
        )
        |> MakePattern

    let CreateInputPatternsForInputPeriod baseIndex inputPeriod (patterns:seq<(int * int) list option>) =
        patterns
        |> Seq.mapi (fun index pattern -> baseIndex + index, if pattern.IsSome then GetPattern pattern.Value inputPeriod else None)
        |> Map.ofSeq

    // Combine two or more patterns
    // Generate a list of input patterns composed of all combinations of the specified patterns taken in sets of two
    // i.e. for each combination of two patterns, take these two patterns and generate a new merged pattern
    let CombinePatterns (existingPatterns:(int * int) list option list) =
        let Merge pattern1Index pattern2Index =
            let pattern1, pattern2 = existingPatterns.[pattern1Index].Value, existingPatterns.[pattern2Index].Value
            let mergedSet =
                let GetNeuronSet index s =
                    let temp = Set.map (fun (time, nindex) -> nindex) s
                    // Warn if any input neurons are shared within the pattern
                    if temp.Count < s.Count then
                        let repeatedNeurons =
                            s
                            |> Seq.fold (fun (state:Map<int, _>) (time, nindex) ->
                                    let count =
                                        if state.ContainsKey nindex then
                                            state.[nindex]
                                        else
                                            0
                                    Map.add nindex (count + 1) state
                                ) Map.empty
                            |> Map.toSeq
                            |> Seq.filter (fun (nindex, count) -> count > 1)
                            |> StringifyWithSeparator " "
                        printfn "Warning: Repeated neurons within pattern %d (%s)" index repeatedNeurons
                    temp

                // Warn if any input neurons are shared across the patterns
                let CheckSharedNeuronsAcrossPatterns neuronSet1 neuronSet2 =
                    let overlap = Set.intersect neuronSet1 neuronSet2
                    if overlap.Count > 0 then
                        printfn "Warning: Patterns %d and %d share %d neurons (%s)" pattern1Index pattern2Index overlap.Count (StringifyWithSeparator " " overlap)

                let set1, set2 = (Set.ofSeq pattern1), (Set.ofSeq pattern2)
                let neuronSet1, neuronSet2 = (GetNeuronSet 0 set1), (GetNeuronSet 1 set2)
                CheckSharedNeuronsAcrossPatterns neuronSet1 neuronSet2
                Set.union set1 set2
            if mergedSet.Count < pattern1.Length + pattern2.Length then
                printfn "Warning: Merged patterns share one or more firing events"
                printfn "Pattern 1: %s" (StringifyWithSeparator " " pattern1)
                printfn "Pattern 2: %s" (StringifyWithSeparator " " pattern2)
            Some(Set.toList mergedSet)

        let GetCombinations numPatterns =
            seq {
                for i in 0..numPatterns-1 do
                    for j in i..numPatterns-1 do
                        if not (i = j) then
                            yield i, j
            }
        let combinedPatterns =
            GetCombinations existingPatterns.Length
            |> Seq.fold (fun state (pattern1, pattern2) ->
                    if existingPatterns.[pattern1].IsSome && existingPatterns.[pattern2].IsSome then
                        let merged = Merge pattern1 pattern2
                        merged :: state
                    else
                        state
                ) []

        combinedPatterns

    let ReadInputPatternsFile path =
        File.ReadLines(path)
        |> Seq.map (fun line -> Some(IOUtilities.ParseStringToList<int * int> line))
        |> Seq.toList

    // Create random patterns and return a list of these patterns plus any additional patterns supplied as arguments
    let CreateRandomInputPatterns numRandomPatterns (additionalPatterns:seq<(int * int) list option>) patternSize maxDelay totalNeurons =

        let CreateInputPatterns numPatterns size maxDelay totalNeurons =
            let CreateRandomInputPattern (r:System.Random) size maxDelay totalNeurons =
                [ for i in 1..size do yield r.Next(1, maxDelay+1), r.Next(totalNeurons) ]
            let r = new System.Random()
            Seq.init numPatterns (fun _ ->
                    Some(CreateRandomInputPattern r size maxDelay totalNeurons)
                )
            |> Seq.toList

        let randomPatterns = CreateInputPatterns numRandomPatterns patternSize maxDelay totalNeurons

        additionalPatterns
        |> Seq.fold (fun state patternData -> patternData :: state) randomPatterns

    // Parse the PNG linkage data in the specified line and combine with the specified Pattern to produce a list of linked events
    let ParseLinkageData template line =

        // create a map of FiringEvents keyed by neuron index
        let CreateNeuronEventMap (template:Pattern) : Map<int, FiringEvent list> =
            template.EventList
            |> Seq.fold (fun (state:Map<int, FiringEvent list>) event ->
                    let key = event.NeuronIndex
                    let eventList =
                        if state.ContainsKey key then
                            state.[key]
                        else
                            []
                    Map.add key (event :: eventList) state
                ) Map.empty

        let mapNeuronEvent = CreateNeuronEventMap template
        let linkedEvents =
            IOUtilities.ParseStringToList<int * int * int * int>(line)
            |> Seq.map (fun (preNeuron, postNeuron, delay, layer) -> preNeuron, postNeuron, delay)  // ignore layer
            |> Seq.distinct                                                                         // ignore duplicates
            |> Seq.map (fun (preNeuron, postNeuron, delay) ->    // See GroupLink
                    let preEvents = mapNeuronEvent.[preNeuron]
                    let postEvents = mapNeuronEvent.[postNeuron]
                    seq {
                        for preEvent in preEvents do
                            for postEvent in postEvents do
                                let eventDelay = postEvent.Time - preEvent.Time
                                // eventDelay range typically 3 - 24; typically larger than delay by 1 - 7
                                if eventDelay > 0 && eventDelay <= 30 then
                                    yield ((preEvent.Time, preEvent.NeuronIndex), (postEvent.Time, postEvent.NeuronIndex))
                    }
                )
            |> Seq.concat
            //|> Seq.cast<SpikingAnalyticsLib.LinkedEventPair>      // Produce a LinkedEventPair (see GeneralTypes)
            |> Seq.map (fun ((t1, n1), (t2, n2)) ->
                let preEvent = new FiringEvent(t1, n1, EventLabel.Foreground)
                let postEvent = new FiringEvent(t2, n2, EventLabel.Foreground)
                new LinkedEventPair(preEvent, postEvent, 0.0, 1.0)
            )
            |> Seq.toList

        linkedEvents

    // Get pattern templates grouped by time (minutes) from a 'groups' file
    // i.e. a file that stores PNG data collected via a PNG search algorithm
    // File format:
    // key      e.g. 6 1        // minutes, totalNumberOfGroups
    // events   (list of FiringEvents in the PNG)
    // ----------------------------------------------------------------------
    // links    (list of event linkage data)
    // ======================================================================
    let PreparePNGTemplatesDictionary groupsFilePath (maxTemplateSize:int option) =
        File.ReadLines(groupsFilePath)
        |> GroupLines 5
        |> Seq.fold (fun (state:Map<int, (int * Pattern) list>) (index, groupedLines) ->
                let key, _ = IOUtilities.ParseStringToTuple<int * int>(groupedLines.[0])
                let template = Pattern.ParseFromString(groupedLines.[1], maxTemplateSize, true)
                let templateList =
                    if state.ContainsKey key then
                        state.[key]
                    else
                        []
                Map.add key ((index, template) :: templateList) state
            ) Map.empty

    // Generate a list of templates of the specified minimum size.
    // Return a TemplateList (a list of Patterns)
    let GetTemplates (templateMap:Map<int, (int * Pattern) list>) minSize =
        let templates =
            templateMap
            |> Map.toSeq
            |> Seq.map (fun (key, keyedTemplates) -> Seq.ofList keyedTemplates)
            |> Seq.concat
            |> Seq.toList
            |> Seq.map (fun (index, template) -> template)
            |> Seq.filter (fun template -> template.Count >= minSize)
            |> Seq.toArray
        new TemplateList(templates)

    // Generate a map of source file indices for each Pattern, keyed by the hash of the Pattern
    // Create a list of templates and source file indices keyed by Pattern hash
    // Multiple templates can have the same hash (in most cases these will be the same template)
    let GetTemplateHashMap (templateMap:Map<int, (int * Pattern) list>) =
        let hashMap =
            templateMap
            |> Map.toSeq
            |> Seq.map (fun (key, keyedTemplates) -> Seq.ofList keyedTemplates)
            |> Seq.concat
            |> Seq.toList
            |> Seq.fold (fun (state:Map<int, (int * Pattern) list>) (sourceIndex, template) ->
                    let hash = template.UniqueID
                    let indexList =
                        if state.ContainsKey hash then
                            state.[hash]
                        else
                            []
                    Map.add hash ((sourceIndex, template) :: indexList) state
                ) Map.empty
        if hashMap.Count = 1 && hashMap.ContainsKey 0 then
            failwith "UniqueID not set on the Patterns in this map"
        else
            hashMap

    /// Create a map of templateID to groups file source index
    // Where templateID is an index into a TemplateList and source index counts the number of 'groups' into the groups file
    // Use for match files with no UniqueID (pattern hash)
    let GetTemplateIDToSourceIndexMap groupsFilePath =
        let templateMap = PreparePNGTemplatesDictionary groupsFilePath None
        let templates =
            let allTemplates = GetTemplates templateMap 0
            allTemplates.DistinctPatterns
        let hashMap = GetTemplateHashMap templateMap
        templates.Patterns
        |> Seq.mapi (fun index template ->
                let hash = template.ComputeHashCode()
                if hashMap.ContainsKey hash then
                    let matches =
                        hashMap.[hash]
                        |> Seq.map (fun (sourceIndex, sourcePattern) ->
                                let matchPattern = sourcePattern :> System.Collections.IStructuralEquatable
                                if matchPattern.Equals(template, System.Collections.Generic.EqualityComparer<FiringEvent>.Default) then
                                    Some(sourceIndex)
                                else
                                    None
                            )
                        |> Seq.choose id
                        |> Seq.toList
                    if matches.Length > 0 then
                        index, matches.Head
                    else
                        failwith "No matches for template ID: %d" index
                else
                    failwith "Key not found for template ID: %d" index
            )
        |> Map.ofSeq

    // Extract all link data from the specified groups file
    let GetAllLinkData groupsFilePath =
        File.ReadLines(groupsFilePath)
        |> GroupLines 5
        |> Seq.map (fun (index, groupedLines) ->
                let parsedTemplate = Pattern.ParseFromString(groupedLines.[1], None, false)
                ParseLinkageData parsedTemplate groupedLines.[3]
            )

    // Extract link data from the specified groups file for the specified list of pattern IDs and Patterns
    let GetIndexedPNGTemplateLinkData groupsFilePath (matchPatterns:Map<int, Pattern>) =
        let sourceIndexSet =
            matchPatterns
            |> Map.toSeq
            |> Seq.map (fun (sourceIndex, pattern) -> sourceIndex)
            |> Set.ofSeq
        File.ReadLines(groupsFilePath)
        |> GroupLines 5
        |> Seq.filter (fun (index, groupedLines) -> sourceIndexSet.Contains index)
        |> Seq.map (fun (index, groupedLines) ->
                let ConfirmTemplateMatch (parsedTemplate:Pattern) (matchPatterns:Map<int, Pattern>) lineIndex =
                    let matchPattern = matchPatterns.[lineIndex]
                    Seq.forall2 (fun evt1 evt2 -> evt1.Equals(evt2)) matchPattern.EventList parsedTemplate.EventList

                let parsedTemplate = Pattern.ParseFromString(groupedLines.[1], None, false)
                if ConfirmTemplateMatch parsedTemplate matchPatterns index then
                    let linkageData = ParseLinkageData parsedTemplate groupedLines.[3]
                    Some(linkageData)
                else
                    printfn "WARNING: failed to match a pattern with the corresponding source at line index %d" index
                    None
            )
        |> Seq.choose id

    // Get the base data from a matches file
    // Returns a map of maps keyed by sample time and then templateID
    // File format:
    // sampleTime, templateID, index, offset, matchCount, matchPattern
    // [match template]
    let GetMatchData matchesFilePath =
        File.ReadLines(matchesFilePath)
        |> GroupLines 2
        |> Seq.fold (fun (state:Map<int, Map<int, (int * int * int * string) list>>) (index, groupedLines) ->
                let sampleTime, templateID, index, offset, matchCount, matchPattern = IOUtilities.ParseStringToTuple<int * int * int * int * int * string>(groupedLines.[0])
                let updatedTemplateMap =
                    let templateMap =       // map keyed by templateID
                        if state.ContainsKey sampleTime then
                            state.[sampleTime]
                        else
                            Map.empty
                    let templateList =      // list of templates with the same templateID
                        if templateMap.ContainsKey templateID then
                            templateMap.[templateID]
                        else
                            []
                    Map.add templateID ((index, offset, matchCount, matchPattern) :: templateList) templateMap
                Map.add sampleTime updatedTemplateMap state
            ) Map.empty

    // Get the base data lines from a matches file
    let GetMatchBaseDataLines matchesFilePath =
        File.ReadLines(matchesFilePath)
        |> GroupLines 2
        |> Seq.map (fun (index, groupedLines) -> groupedLines.[0])

    // Get all the matching templates at the specified sample time
    // Returns a map keyed by templateID
    // File format:
    // sampleTime, templateID, index, offset, matchCount, matchPattern
    // [match template]
    let GetMatchingTemplatesAtTime_V1 sampleTime matchesFilePath initialMatchLength =
        let matchTime = sprintf "%d " sampleTime
        let initialMatchPattern = String.replicate initialMatchLength "1"
        File.ReadLines(matchesFilePath)
        |> GroupLines 2
        |> Seq.filter (fun (index, groupedLines) -> groupedLines.[0].StartsWith(matchTime))
        |> Seq.fold (fun (templateMap:Map<int, (int * int * int * string * Pattern) list>) (index, groupedLines) ->
                let headerLine = groupedLines.[0]
                let time, templateID, index, offset, matchCount, matchPattern = IOUtilities.ParseStringToTuple<int * int * int * int * int * string>(headerLine)
                if time <> sampleTime then
                    failwith "Sample time matching error"
                if initialMatchLength = 0 || matchPattern.StartsWith(initialMatchPattern) then
                    let templateList =      // list of templates with the same templateID
                        if templateMap.ContainsKey templateID then
                            templateMap.[templateID]
                        else
                            []
                    let template = Pattern.ParseFromString(groupedLines.[1], None, false)
                    Map.add templateID ((index, offset, matchCount, matchPattern, template) :: templateList) templateMap
                else
                    templateMap
            ) Map.empty

    // Get all the matching templates at the specified sample time
    // Returns a map keyed by templateID and uniqueID (pattern hash)
    // File format:
    // sampleTime, templateID uniqueID, index, offset, matchCount, matchPattern
    // [match template]
    let GetMatchingTemplatesAtTime_V2 sampleTime matchesFilePath initialMatchLength =
        let matchTime = sprintf "%d " sampleTime
        let initialMatchPattern = String.replicate initialMatchLength "1"
        File.ReadLines(matchesFilePath)
        |> GroupLines 2
        |> Seq.filter (fun (index, groupedLines) -> groupedLines.[0].StartsWith(matchTime))
        |> Seq.fold (fun (templateMap:Map<int * int, (int * int * int * string * Pattern) list>) (index, groupedLines) ->
                let headerLine = groupedLines.[0]
                let time, templateID, uniqueID, index, offset, matchCount, matchPattern = IOUtilities.ParseStringToTuple<int * int * int * int * int * int * string>(headerLine)
                if time <> sampleTime then
                    failwith "Sample time matching error"
                if initialMatchLength = 0 || matchPattern.StartsWith(initialMatchPattern) then
                    let templateList =      // list of templates with the same templateID
                        if templateMap.ContainsKey (templateID, uniqueID) then
                            templateMap.[templateID, uniqueID]
                        else
                            []
                    let template = Pattern.ParseFromString(groupedLines.[1], None, false)
                    Map.add (templateID, uniqueID) ((index, offset, matchCount, matchPattern, template) :: templateList) templateMap
                else
                    templateMap
            ) Map.empty

    // Get the initial triplet from each pattern in a file of unique templates
    let GetInitialTriplets templatesFilePath =
        File.ReadLines(templatesFilePath)
        |> Seq.map (fun line ->
                let template = Pattern.ParseFromString(line, None, false)
                template.InitialTriplet
            )

    let SavePatterns (inputPatterns:(int * int) list option list) path =
        let lines =
            inputPatterns
            |> Seq.mapi (fun index pattern ->
                    if pattern.IsSome then
                        StringifyWithSeparator " " pattern.Value
                    else
                        "[]"
                )
        File.WriteAllLines(path, lines)
